from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_community.chat_models import ChatTongyi
from typing import List 


class Entities(BaseModel):
    """Identifying information about entities."""
    names: List[str] = Field(
        ...,
        description="All the person, organization, or business entities that "
        "appear in the text",
    )


def get_prompt():
    return ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "You are extracting organization and person entities from the text.",
            ),
            (
                "human",
                "Use the given format to extract information from the following "
                "input: {question}",
            ),
        ]
    )


def get_entity_chain():
    llm = ChatTongyi(api_key='sk-780da555f1714f4b81c1bfc553cfd74a')
    return get_prompt() | llm.with_structured_output(Entities)


if __name__ == "__main__":
    entity_chain = get_entity_chain() 
    result = entity_chain.invoke({"question" : "姚明和叶莉的关系是?"}).names
    print(result)
